import numpy as np
import os
import json
from tqdm import tqdm


base_dir = '../json/'
model1 = 'resnet101_ibn_a'  # 55.58 42.67
model2 = 'resnet101_ibn_b'  # 54.8
model3 = 'se_resnet101_ibn_a'  # 55.4 41.09
model4 = 'hrnet/normal'  # 878 931

model5 = 'resnet50_ibn_a/pseudo'  # 55.79
model6 = 'efficientnet-b2/pseudo'
model7 = 'regnet/pseudo'  # 54.95
model8 = 'hrnet/pseudo'

query = np.load(base_dir + model1 + '/query_path_2.npy')
gallery = np.load(base_dir + model1 + '/gallery_path_2.npy')
print(len(query), len(gallery))

distmat = np.load(base_dir + model1 + '/distmat_2.npy')  # 各个模型的距离矩阵相加
distmat += np.load(base_dir + model2 + '/distmat_2.npy')
distmat += np.load(base_dir + model3 + '/distmat_2.npy')
# distmat += np.load(base_dir + model4 + '/distmat_t.npy')

# distmat = np.load(base_dir + model5 + '/distmat_t.npy')
# distmat += np.load(base_dir + model6 + '/distmat_t.npy')
# distmat += np.load(base_dir + model7 + '/distmat_t.npy')
# distmat += np.load(base_dir + model8 + '/distmat_t.npy')

# base_dir = '../save_model/single/'
# model1 = '12-RMGL+Douzi+LH-RSA-noRP-resnet101-ibn-a-avg+max'  # 49.55
# model2 = '13-se_resnet101'  #47.82
# # model3 = '8-RMGL+Douzi-LH-RSA-noRP-avg' #47.69
# model4 = '14-resnet-ibn-b'  # 47.22
# query = np.load(base_dir + model1 + '/query_path_2.npy')
# gallery = np.load(base_dir + model1 + '/gallery_path_2.npy')
#
# distmat = np.load(base_dir + model1 + '/distmat_2.npy')  # 各个模型的距离矩阵相加
# distmat += np.load(base_dir + model2 + '/distmat_2.npy')
# # distmat += np.load(base_dir + model3 + '/distmat_2.npy')
# distmat += np.load(base_dir + model4 + '/distmat_2.npy')

print(distmat.shape)  # (querry,gallery)距离矩阵
indexes = np.argsort(distmat, axis=1)  # (querry,gallery)距离最近的索引矩阵

res = {}
for idx, index in tqdm(enumerate(indexes)):
    query_ = os.path.basename(query[idx])
    gallery_ = [os.path.basename(i) for i in gallery[index][:200].tolist()]
    res[query_] = gallery_


data = dict()
for k, v in res.items():
    data[k] = v

save_path = 'submit_final.json'
print("Writing to {}".format(save_path))
json.dump(data, open(save_path, 'w'))
